#!/usr/bin/env python
# -*- coding: utf-8 -*-
import pathlib
import time
import pandas as pd
from absl import app, flags, logging
from futu import OpenQuoteContext, KLType, AuType, RET_OK, Market, SecurityType

flags.DEFINE_string('cache_dir', '/data/quant/hk', 'Cache directory')
flags.DEFINE_string('start_date', None, 'Data start date')
flags.DEFINE_string('end_date', None, 'Data end date')
flags.DEFINE_list('codes', [], 'Codes to read')
flags.DEFINE_enum('freq', '1d', ['1d', '1m'], 'k bar frequency')

FLAGS = flags.FLAGS


def refresh_bars(quote_ctx: OpenQuoteContext, code: str, start: str, end: str,
                 kltype: KLType, ins):
    bars = []
    max_count = 1000
    ret, data, page_req_key = quote_ctx.request_history_kline(
        code,
        start=start,
        end=end,
        ktype=kltype,
        autype=AuType.QFQ,
        max_count=max_count)  # 请求第一页
    if ret == RET_OK:
        bars.append(data)
    else:
        logging.error('error:', data)
    while page_req_key is not None:  # 请求后面的所有结果
        logging.info('%s data read' % (len(bars) * max_count))
        ret, data, page_req_key = quote_ctx.request_history_kline(
            code,
            start=start,
            end=end,
            ktype=kltype,
            autype=AuType.QFQ,
            max_count=max_count,
            page_req_key=page_req_key)  # 请求翻页后的数据
        if ret == RET_OK:
            bars.append(data)
        else:
            logging.error('error:', data)
        time.sleep(1)
    logging.info('All pages are finished!')
    df = pd.concat(bars)
    df.rename(columns={
        'code': 'order_book_id',
        'time_key': 'datetime'
    },
        inplace=True,
        errors='raise')
    if ins.stock_type == 'IDX':
        df['volume'] = 10000
    return df


def refresh_basic_info(quote_ctx: OpenQuoteContext):
    dfs = []
    for stype in [SecurityType.STOCK, SecurityType.ETF, SecurityType.IDX]:
        logging.info('Read %s' % stype)
        ret, data = quote_ctx.get_stock_basicinfo(Market.HK, stype)
        if ret == RET_OK:
            df = data.rename(columns={
                'code': 'order_book_id',
                'name': 'symbol',
                'listing_date': 'listed_date',
                'lot_size': 'round_lot',
                'delisted_date': 'de_listed_date',
            },
                errors='raise')
            df['type'] = 'CS'
            df['market_tplus'] = 0
            df['account_type'] = 'STOCK'
            df['trading_hours'] = '09:31-12:00,13:31-16:00'
            df['board_type'] = 'MainBoard'
            df['exchange'] = 'HK'
            dfs.append(df)
        else:
            return ret, data
    return RET_OK, pd.concat(dfs).set_index('order_book_id')


def main(_):
    cache_dir = pathlib.Path(FLAGS.cache_dir).joinpath(FLAGS.end_date)
    assert cache_dir.exists()
    quote_ctx = OpenQuoteContext(host='127.0.0.1', port=11111)
    f_all_instruments = cache_dir.joinpath('all_instruments.csv')
    if f_all_instruments.exists():
        all_instruments = pd.read_csv(f_all_instruments,
                                      index_col=['order_book_id'])
    else:
        logging.info('Read all_instruments...')
        ret, all_instruments = refresh_basic_info(quote_ctx)
        if ret == RET_OK:
            all_instruments.to_csv(f_all_instruments, index=False)
        else:
            logging.error("error %s:" % all_instruments)

    if FLAGS.freq == '1d':
        kltype = KLType.K_DAY
    else:
        kltype = KLType.K_1M
    kdir = cache_dir.joinpath(FLAGS.freq)
    for code in FLAGS.codes:
        kpath = kdir.joinpath('%s.csv' % code)
        ins = all_instruments.loc[code]
        if not kpath.exists():
            logging.info('Read %s' % code)
            df = refresh_bars(quote_ctx, code, FLAGS.start_date,
                              FLAGS.end_date, kltype, ins)
            df.to_csv(kpath, index=False)
    quote_ctx.close()  # 结束后记得关闭当条连接，防止连接条数用尽


if __name__ == '__main__':
    app.run(main)
